import torch
import json
import numpy as np
import pandas as pd

from net import LSTMModel
from utils import prepare_data, Normalizer, loadDict4test, configs, TimeDelta

def predict(dict, predict_steps=3):
	results = []

	all_dic = loadDict4test(dict)
	for data_type, group_info in all_dic.items():
		config = configs[data_type]
		for stationID, data_dic in group_info.items():
			if not data_dic['is_legal']:
				raise ValueError(f"Data for station {stationID} is not legal.") 
			data_df = data_dic['df']

			# normalize
			with open(f'model_save/{data_type}/{stationID}_scaler.json', 'r') as f:
				scaler_info = json.load(f)
			scaler = Normalizer(
				mu=np.array(scaler_info["mu"]),
				sd=np.array(scaler_info["sd"])
			)

			data_df['time'] = pd.to_datetime(data_df['time'])
			freq = TimeDelta[data_type]
			last = data_df['time'].iloc[-1]

			match data_type:
				case 'movement':
					vals = np.stack([
						data_df['val_e'].values,
						data_df['val_n'].values,
						data_df['val_u'].values,
						], axis=-1)
				case _:
					vals = data_df['value'].values

			vals = scaler.test_transform(vals)		

			model = LSTMModel(
				input_size=config["model"]["input_size"],
				hidden_layer_size=config["model"]["lstm_size"],
				num_layers=config["model"]["num_lstm_layers"],
				output_size=config["model"]["output_size"],
				dropout=config["model"]["dropout"],
			)
			model.load_state_dict(torch.load(f"model_save/{data_type}/{stationID}.pt", weights_only=True, map_location=torch.device('cpu')))
			model.eval()

			with torch.no_grad():
				for _ in range(predict_steps):
					next_ts = last + pd.tseries.frequencies.to_offset(freq)
					match data_type:
						case 'movement':
							data_x = torch.from_numpy(vals)[None, ...].to("cpu").float()
						case _:
							data_x = torch.from_numpy(vals)[None, ..., None].to("cpu").float()
					predicted_val = model(data_x)
					match data_type:
						case "movement":
							transformed_prediction = scaler.inverse_transform(predicted_val)
							results.append({
								"modelType": "LSTM",
								"stationID": stationID,
								"stationType":data_type,
								"time": str(next_ts),
								"val_e": transformed_prediction[:, 0].item(),
								"val_n": transformed_prediction[:, 1].item(),
								"val_u": transformed_prediction[:, 2].item(),

							})

						case _:
							results.append({
								"modelType": "LSTM",
								"stationID": stationID,
								"stationType":data_type,
								"time": str(next_ts),
								"value": scaler.inverse_transform(predicted_val).item()
							})

					# update the last timestamp and values
					match data_type:
						case 'movement':
							vals = np.vstack([vals, predicted_val])
							vals = vals[1:]
						case _:
							vals = np.append(vals, predicted_val.squeeze().item())
							vals = vals[1:]
							
					last = next_ts

	return results

if __name__ == "__main__":
	pressure_data = {
		"startTime":"2005-5-8 12:00:00",
		"endTime":"2005-5-8 17:00:00",
		"data":
		[
			{"stationID":"85551", "StationType": "渗压", "time":"2005-5-8 12:00:00","value":44.88},
			{"stationID":"85551", "StationType": "渗压", "time":"2005-5-8 13:00:00","value":45.88},
			{"stationID":"85551", "StationType": "渗压", "time":"2005-5-8 14:00:00","value":47.88},
			{"stationID":"85551", "StationType": "渗压", "time":"2005-5-8 15:00:00","value":49.58},
			{"stationID":"85551", "StationType": "渗压", "time":"2005-5-8 16:00:00","value":43.48},
			{"stationID":"85551", "StationType": "渗压", "time":"2005-5-8 17:00:00","value":42.38},
			
			{"stationID":"85552", "StationType": "渗压", "time":"2005-5-8 12:00:00","value":44.88},
			{"stationID":"85552", "StationType": "渗压", "time":"2005-5-8 13:00:00","value":45.88},
			{"stationID":"85552", "StationType": "渗压", "time":"2005-5-8 14:00:00","value":47.88},
			{"stationID":"85552", "StationType": "渗压", "time":"2005-5-8 15:00:00","value":49.58},
			{"stationID":"85552", "StationType": "渗压", "time":"2005-5-8 16:00:00","value":43.48},
			{"stationID":"85552", "StationType": "渗压", "time":"2005-5-8 17:00:00","value":42.38},
	]}
	results = predict(pressure_data, predict_steps=3)
	print(results)

	seepage_data = {
		"startTime":"2005-5-8 12:00:00",
		"endTime":"2005-5-8 17:00:00",
		"data":
		[
			{"stationID":"85551", "StationType": "渗流", "time":"2005-5-8 12:00:00","value":0.00},
			{"stationID":"85551", "StationType": "渗流", "time":"2005-5-8 13:00:00","value":0.00},
			{"stationID":"85551", "StationType": "渗流", "time":"2005-5-8 14:00:00","value":0.00},
			{"stationID":"85551", "StationType": "渗流", "time":"2005-5-8 15:00:00","value":0.00},
			{"stationID":"85551", "StationType": "渗流", "time":"2005-5-8 16:00:00","value":0.00},
			{"stationID":"85551", "StationType": "渗流", "time":"2005-5-8 17:00:00","value":0.00},
	]}
	results = predict(seepage_data, predict_steps=3)
	print(results)

	radar_data = {
		"startTime":"2005-5-8 12:00:00",
		"endTime":"2005-5-8 17:00:00",
		"data":
		[
			{"stationID":"4209021118", "StationType": "超声波雷达", "time":"2005-5-8 12:00:00","value":-0.19},
			{"stationID":"4209021118", "StationType": "超声波雷达", "time":"2005-5-8 13:00:00","value":-0.2},
			{"stationID":"4209021118", "StationType": "超声波雷达", "time":"2005-5-8 14:00:00","value":-0.16},
			{"stationID":"4209021118", "StationType": "超声波雷达", "time":"2005-5-8 15:00:00","value":0.23},
			{"stationID":"4209021118", "StationType": "超声波雷达", "time":"2005-5-8 16:00:00","value":0.24},
			{"stationID":"4209021118", "StationType": "超声波雷达", "time":"2005-5-8 17:00:00","value":0.03},
			
			{"stationID":"4209027018", "StationType": "超声波雷达", "time":"2005-5-8 12:00:00","value":-0.19},
			{"stationID":"4209027018", "StationType": "超声波雷达", "time":"2005-5-8 13:00:00","value":-0.2},
			{"stationID":"4209027018", "StationType": "超声波雷达", "time":"2005-5-8 14:00:00","value":-0.16},
			{"stationID":"4209027018", "StationType": "超声波雷达", "time":"2005-5-8 15:00:00","value":0.23},
			{"stationID":"4209027018", "StationType": "超声波雷达", "time":"2005-5-8 16:00:00","value":0.24},
			{"stationID":"4209027018", "StationType": "超声波雷达", "time":"2005-5-8 17:00:00","value":0.03},
	]}
	results = predict(radar_data, predict_steps=3)
	print(results)
	

	movement_data = {
		"startTime":"2005-5-8 12:00:00",
		"endTime":"2005-5-8 17:00:00",
		"data":[
			{"stationID":"4209023000101", "StationType": "变形监测", "time":"2005-5-8 12:00:00","val_e":0.1, "val_n": -0.3 , "val_u": 0.1},
			{"stationID":"4209023000101", "StationType": "变形监测", "time":"2005-5-8 13:00:00","val_e":0.5, "val_n": 0.18 , "val_u": 0.2},
			{"stationID":"4209023000101", "StationType": "变形监测", "time":"2005-5-8 14:00:00","val_e":0.8, "val_n": -0.13, "val_u": 0.3},
			{"stationID":"4209023000101", "StationType": "变形监测", "time":"2005-5-8 15:00:00","val_e":1.2, "val_n": -0.12, "val_u": 0.5},
			{"stationID":"4209023000101", "StationType": "变形监测", "time":"2005-5-8 16:00:00","val_e":0.7, "val_n": -0.18, "val_u": -0.1},
			{"stationID":"4209023000101", "StationType": "变形监测", "time":"2005-5-8 17:00:00","val_e":0.7, "val_n": 0.11 , "val_u": -0.7},

			{"stationID":"4209023000102", "StationType": "变形监测", "time":"2005-5-8 12:00:00","val_e":0.1, "val_n": -0.3 , "val_u": 0.1},
			{"stationID":"4209023000102", "StationType": "变形监测", "time":"2005-5-8 13:00:00","val_e":0.5, "val_n": 0.18 , "val_u": 0.2},
			{"stationID":"4209023000102", "StationType": "变形监测", "time":"2005-5-8 14:00:00","val_e":0.8, "val_n": -0.13, "val_u": 0.3},
			{"stationID":"4209023000102", "StationType": "变形监测", "time":"2005-5-8 15:00:00","val_e":1.2, "val_n": -0.12, "val_u": 0.5},
			{"stationID":"4209023000102", "StationType": "变形监测", "time":"2005-5-8 16:00:00","val_e":0.7, "val_n": -0.18, "val_u": -0.1},
			{"stationID":"4209023000102", "StationType": "变形监测", "time":"2005-5-8 17:00:00","val_e":0.7, "val_n": 0.11 , "val_u": -0.7},
		]
	}
	results = predict(movement_data, predict_steps=3)
	print(results)

